import cv2
import base64
import requests
import json
from io import BytesIO
from PIL import Image
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image as RosImage
from cv_bridge import CvBridge
from tf2_ros import TransformException
from tf2_ros.buffer import Buffer
from tf2_ros.transform_listener import TransformListener
from transforms3d.euler import quat2euler
import math
import threading
import time
import csv
from datetime import datetime

class ImageAnalyzer(Node):
    def __init__(self):
        super().__init__('image_analyzer')
        self.bridge = CvBridge()
        self.subscription = self.create_subscription(
            RosImage,
            '/fishbot_camera_raw',
            self.image_callback,
            10)
        self.last_process_time = self.get_clock().now()
        self.process_interval = 1.0
        
        # 初始化 tf 相关内容
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)
        
        # 初始化线程相关变量
        self.pose_lock = threading.Lock()
        self.current_pose = None
        self.running = True
        
        # 启动独立线程处理坐标获取
        self.pose_thread = threading.Thread(target=self.pose_update_loop, daemon=True)
        self.pose_thread.start()
        
        # 添加CSV文件相关初始化
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.csv_filename = f'scene_data_{timestamp}.csv'
        with open(self.csv_filename, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['描述', 'X坐标', 'Y坐标', '航向角'])

    def pose_update_loop(self):
        """独立线程循环获取位置信息"""
        while self.running and rclpy.ok():
            try:
                transform = self.tf_buffer.lookup_transform(
                    'odom',
                    'base_footprint',
                    rclpy.time.Time())
                
                q = transform.transform.rotation
                roll, pitch, yaw = quat2euler([q.w, q.x, q.y, q.z])
                yaw_deg = math.degrees(yaw)
                if yaw_deg < 0:
                    yaw_deg += 360
                
                with self.pose_lock:
                    self.current_pose = (
                        transform.transform.translation.x,
                        transform.transform.translation.y,
                        transform.transform.translation.z,
                        yaw_deg,
                        'odom'
                    )
            except Exception:
                pass
            
            time.sleep(0.1)

    def image_callback(self, msg):
        try:
            self.current_frame = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            cv2.imshow('Camera', self.current_frame)
            cv2.waitKey(1)
            
            current_time = self.get_clock().now()
            if (current_time - self.last_process_time).nanoseconds / 1e9 >= self.process_interval:
                description = self.analyze_image(self.current_frame)
                with self.pose_lock:
                    current_pose = self.current_pose
                
                if current_pose:
                    x, y, z, yaw, frame = current_pose
                    location_info = f"\n机器人位置 ({frame}坐标系): x={x:.2f}, y={y:.2f}, z={z:.2f}, 航向角={yaw:.2f}°"
                    print("画面分析结果：", description + location_info)
                    
                    # 添加CSV记录
                    with open(self.csv_filename, 'a', newline='', encoding='utf-8') as f:
                        writer = csv.writer(f)
                        writer.writerow([description.strip(), f"{x:.2f}", f"{y:.2f}", f"{yaw:.2f}"])
                else:
                    print("画面分析结果：", description)
                self.last_process_time = current_time
                
        except Exception as e:
            self.get_logger().error(f"处理出错: {str(e)}")

    def encode_image_to_base64(self, image):
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        pil_image = Image.fromarray(image_rgb)
        buffered = BytesIO()
        pil_image.save(buffered, format="JPEG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')

    def analyze_image(self, frame):
        try:
            base64_image = self.encode_image_to_base64(frame)
            data = {
                "model": "minicpm-v",
                "messages": [
                    {
                        "role": "user",
                        "content": """
                        分析这个画面。以方便产生场景记忆点的方式提取关键特征。比如地标性的景观，特定的物体，特定的装饰。最后输出格式：
                        特征点：描述
                        特征点：描述
                        """,
                        "images": [base64_image]
                    }
                ],
                "stream": False  # 关流式响应
            }
            
            response = requests.post('http://localhost:11434/api/chat', json=data)
            
            # 读取整个响应文本
            response_text = response.text.strip()
            # 获取最后一个JSON对象（完整的响应）
            last_response = json.loads(response_text.split('\n')[-1])
            
            if not response.ok:
                return f"API请求失败: {response.status_code}"
                
            if 'error' in last_response:
                return f"API返回错误: {last_response['error']}"
                
            if 'message' in last_response and 'content' in last_response['message']:
                return last_response['message']['content']
            else:
                print("API完整响应:", last_response)
                return "API返回格式异常"
                
        except requests.exceptions.RequestException as e:
            return f"网络请求错误: {str(e)}"
        except json.JSONDecodeError as e:
            return f"JSON解析错误: {str(e)}"
        except Exception as e:
            return f"未知错误: {str(e)}"

    def __del__(self):
        self.running = False
        if hasattr(self, 'pose_thread'):
            self.pose_thread.join(timeout=1.0)

def main():
    rclpy.init()
    analyzer = ImageAnalyzer()
    try:
        rclpy.spin(analyzer)
    except KeyboardInterrupt:
        pass
    finally:
        # 清理资源
        analyzer.running = False
        if hasattr(analyzer, 'pose_thread'):
            analyzer.pose_thread.join(timeout=1.0)
        cv2.destroyAllWindows()
        analyzer.destroy_node()
        try:
            rclpy.shutdown()
        except Exception as e:
            pass  # 忽略关闭时的错误

if __name__ == '__main__':
    main()
